import tushare as ts
import pandas as pd
import numpy as np


class GetHist:
    def __init__(self, code, start, end):
        data_ = ts.get_hist_data(code, start=start, end=end)
        self.code_data = data_

    def get_hist_numpy(self):
        df = self.code_data
        data_ = np.array(df.T)
        return data_

    def get_hist_global(self, data_):
        open_price = data_[0, -1]
        close_price = data_[2, 0]
        return float(close_price) - float(open_price)

    def fall_one(self, data_, step, expect=True):
        i = 0
        res = []
        while True:
            tmp = data_[:, i:i + step]
            tmp = list(np.squeeze(tmp.reshape([-1, 1])))
            if expect:
                exp = self.get_hist_global(data_[:, i + step: i + step * 2])
                tmp.append(exp)
            res.append(tmp)
            if i + step >= data_.shape[1] or (expect and i + step * 2 >= data_.shape[1]):
                break
            i += 1
        return np.array(res).T


if __name__ == '__main__':
    gh = GetHist('600408', '2017-01-01', '2017-01-10')
    gh = GetHist('600408', '2017-03-01', '2017-03-30')

    data = gh.get_hist_numpy()
    data = gh.fall_one(data, 2, True)

    # print(data)
